from datetime import datetime
import logging

import pandas as pd

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s | %(filename)s - {%(lineno)d} - [%(levelname)s] - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S', encoding='utf-8')


class StockStrategy:

    def __init__(self, data_list, p1=1, p2=5, p3=1):
        if data_list is not None and len(data_list) > 0:
            self.data_list = data_list
            self.df = pd.DataFrame(self.data_list,
                                   columns=["code", "name", "date", "open", "close", "high", "low", "vol", "pchange",
                                            "zl2",
                                            "zl1", "ma5", "ma10", "ma20", "ma30", "ma60", "v5", "v10", "v20", "dif",
                                            "dea",
                                            "macd", "kdj_k", "kdj_d", "kdj_j", "rsi_6", "atr_14"])
            self.p1 = p1
            self.p2 = p2
            self.p3 = p3
            ## fillna
            self.df["zl1"] = self.df["zl1"].fillna(0)
            self.df["zl2"] = self.df["zl2"].fillna(0)
            self.df["ma5"] = self.df["ma5"].fillna(0)
            self.df["ma10"] = self.df["ma10"].fillna(0)
            self.df["ma20"] = self.df["ma20"].fillna(0)
            self.df["ma30"] = self.df["ma30"].fillna(0)
            self.df["ma60"] = self.df["ma60"].fillna(0)
            self.df["v5"] = self.df["v5"].fillna(0)
            self.df["v10"] = self.df["v10"].fillna(0)
            self.df["v20"] = self.df["v20"].fillna(0)
            self.df["dif"] = self.df["dif"].fillna(0)
            self.df["dea"] = self.df["dea"].fillna(0)
            self.df["macd"] = self.df["macd"].fillna(0)
            self.df["kdj_k"] = self.df["kdj_k"].fillna(0)
            self.df["kdj_d"] = self.df["kdj_d"].fillna(0)
            self.df["kdj_j"] = self.df["kdj_j"].fillna(0)
            self.df["rsi_6"] = self.df["rsi_6"].fillna(0)
            self.df["atr_14"] = self.df["atr_14"].fillna(0)
            ## data type change
            self.df["vol"] = self.df["vol"].astype(float)
        else:
            logging.error("data_list is None or empty")
            raise ValueError("data_list is None or empty")

    def strategy(self, dayTime):
        current_day = datetime.strptime(dayTime, "%Y-%m-%d").date()
        current_df = self.df[self.df["date"] == current_day].copy()

        if current_df.size == 0:
            return

        current_df_index = current_df.index[0]

        """
        C1 := ZL2 > ZL1 AND ZL2 > REF(ZL2, 1);
        """
        if current_df.loc[current_df_index, "zl2"] > current_df.loc[current_df_index, "zl1"] and current_df.loc[
            current_df_index, "zl2"] > self.df.loc[current_df_index - 1, "zl2"]:
            current_df.loc[current_df_index, "c1"] = True
        else:
            current_df.loc[current_df_index, "c1"] = False

        """
        C2 := LAST(C <= ZL2, P2, P2 - P1);
        """
        flag1 = True
        for i in range(self.p2 + 1):
            if i >= self.p2 - self.p1:
                if self.df.loc[current_df_index - i, "close"] > self.df.loc[current_df_index - i, "zl2"]:
                    flag1 = False
                    break
        if flag1:
            current_df.loc[current_df_index, "c2"] = True
        else:
            current_df.loc[current_df_index, "c2"] = False

        """
        C3 := EVERY(C > ZL2, P1);
        """
        flag2 = True
        for i in range(self.p1):
            if self.df.loc[current_df_index - i, "close"] <= self.df.loc[current_df_index - i, "zl2"]:
                flag2 = False
                break
        if flag2:
            current_df.loc[current_df_index, "c3"] = True
        else:
            current_df.loc[current_df_index, "c3"] = False

        """
        V_MA5:=MA(V,5);
        VOL_RATIO:=V/MA(V,20);
        C4_1:=V>V_MA5*0.9 AND V<V_MA5*1.5;
        C4_2:=VOL_RATIO > 1.2;
        C4:=EVERY(C4_1,P1) AND C4_2;
        """
        if current_df.loc[current_df_index, "v5"] * 0.9 < current_df.loc[current_df_index, "vol"] < current_df.loc[
            current_df_index, "v5"] * 1.5:
            current_df.loc[current_df_index, "c4_1"] = True
        else:
            current_df.loc[current_df_index, "c4_1"] = False
        if current_df.loc[current_df_index, "vol"] / current_df.loc[current_df_index, "v20"] > 1.2:
            current_df.loc[current_df_index, "c4_2"] = True
        else:
            current_df.loc[current_df_index, "c4_2"] = False
        flag3 = True
        for i in range(self.p1):
            if self.df.loc[current_df_index - i, "v5"] * 0.9 < self.df.loc[current_df_index - i, "vol"] < \
                    self.df.loc[current_df_index - i, "v5"] * 1.5:
                flag3 = flag3 and True
            else:
                flag3 = flag3 and False
                break
        current_df.loc[current_df_index, "c4_1"] = flag3
        current_df.loc[current_df_index, "c4"] = current_df.loc[current_df_index, "c4_1"] and current_df.loc[
            current_df_index, "c4_2"]

        """
        C5:=NOT(INBLOCK('ST版块'));
        """
        current_df.loc[current_df_index, "c5"] = current_df.loc[current_df_index, "name"].find("ST") == -1

        """
        MACD_DIF:=EMA(C,12)-EMA(C,26);
        MACD_DEA:=EMA(MACD_DIF,9);
        C6:=MACD_DIF>MACD_DEA AND MACD_DIF>REF(MACD_DIF,1) AND MACD_DIF >0 AND MACD_DEA >0;
        """
        current_df.loc[current_df_index, "c6"] = current_df.loc[current_df_index, "dif"] > current_df.loc[
            current_df_index, "dea"] and \
                                                 current_df.loc[current_df_index, "dif"] > self.df.loc[
                                                     current_df_index - 1, "dif"] and \
                                                 current_df.loc[current_df_index, "dif"] > 0 and current_df.loc[
                                                     current_df_index, "dea"] > 0

        """
        C7:=C>MAX(ZL2,O) AND C>MA(C,20);
        """
        current_df.loc[current_df_index, "c7"] = current_df.loc[current_df_index, "close"] > max(
            current_df.loc[current_df_index, "zl2"],
            current_df.loc[current_df_index, "open"]) and \
                                                 current_df.loc[current_df_index, "close"] > \
                                                 current_df.loc[current_df_index, "ma20"]

        """
        C8:=(H-C)/C<=P3/100 AND (C-L)/C>0.015;
        """
        current_df.loc[current_df_index, "c8"] = ((current_df.loc[current_df_index, "high"] - current_df.loc[
            current_df_index, "close"]) /
                                                  current_df.loc[current_df_index, "close"] <= self.p3 / 100) and (
                                                         (current_df.loc[current_df_index, "close"] - current_df.loc[
                                                             current_df_index, "low"]) / current_df.loc[
                                                             current_df_index, "close"]) > 0.015

        """
        C9:=RSI(6)>50 AND KDJ.J>50;
        """
        current_df.loc[current_df_index, "c9"] = current_df.loc[current_df_index, "rsi_6"] > 50 and current_df.loc[
            current_df_index, "kdj_j"] > 50

        """
        C10:=CLOSE/LLV(LOW,20)>1.1;
        """
        current_df.loc[current_df_index, "c10"] = current_df.loc[current_df_index, "close"] / (
            self.df.loc[current_df_index - 20:current_df_index, "low"].min()) > 1.1

        """
        MACD_TOP := C > REF(C,1) AND MACD_DIF < REF(MACD_DIF,1);
        C11:=NOT(MACD_TOP);
        """
        current_df.loc[current_df_index, "macd_top"] = current_df.loc[current_df_index, "close"] > self.df.loc[
            current_df_index - 1, "close"] and current_df.loc[current_df_index, "dif"] < self.df.loc[
                                                           current_df_index - 1, "dif"]
        current_df.loc[current_df_index, "c11"] = not current_df.loc[current_df_index, "macd_top"]

        """
        ATR14:=ATR(14);
        C12:=C/REF(C,1) < 1+ATR14/C*2;
        """
        current_df.loc[current_df_index, "c12"] = (current_df.loc[current_df_index, "close"] / self.df.loc[current_df_index - 1, "close"]) < 1 + current_df.loc[current_df_index, "atr_14"] / current_df.loc[current_df_index, "close"] * 2

        """
        BUY_SIGNAL:=CROSS(ZL2,ZL1) OR (C>REF(HHV(H,5),1));
        """

        df = current_df[(current_df["c1"] == True) & (current_df["c2"] == True) & (current_df["c3"] == True) & (current_df["c4"] == True) & (current_df["c5"] == True) & (current_df["c6"] == True) & (current_df["c7"] == True) & (current_df["c8"] == True) & (current_df["c9"] == True) & (current_df["c10"]== True) & (current_df["c11"] == True) & (current_df["c12"] == True)]
        if df is not None and df.shape[0] > 0:
            return dayTime
        else:
            return None

    def update(self):
        if self.df is not None and self.df.size > 0:
            return self.df
        else:
            return None
